/**
 * Writed by yura and mak
 */

package wz.legend;


import org.apache.commons.digester.Digester;
import org.apache.http.ConnectionReuseStrategy;
import org.apache.http.HttpEntity;
import org.apache.http.HttpEntityEnclosingRequest;
import org.apache.http.HttpException;
import org.apache.http.HttpHost;
import org.apache.http.HttpRequest;
import org.apache.http.HttpResponse;
import org.apache.http.HttpStatus;
import org.apache.http.HttpVersion;
import org.apache.http.ProtocolVersion;
import org.apache.http.impl.DefaultConnectionReuseStrategy;
import org.apache.http.impl.DefaultHttpClientConnection;
import org.apache.http.impl.DefaultHttpRequestFactory;
import org.apache.http.impl.DefaultHttpResponseFactory;
import org.apache.http.impl.DefaultHttpServerConnection;
import org.apache.http.impl.SocketHttpServerConnection;
import org.apache.http.impl.io.HttpRequestParser;
import org.apache.http.impl.io.HttpResponseParser;
import org.apache.http.impl.io.HttpResponseWriter;
import org.apache.http.impl.io.SocketInputBuffer;
import org.apache.http.impl.io.SocketOutputBuffer;
import org.apache.http.io.HttpMessageWriter;
import org.apache.http.io.SessionOutputBuffer;
import org.apache.http.message.BasicHttpEntityEnclosingRequest;
import org.apache.http.message.BasicLineFormatter;
import org.apache.http.message.BasicLineParser;
import org.apache.http.message.LineFormatter;
import org.apache.http.params.BasicHttpParams;
import org.apache.http.params.DefaultedHttpParams;
import org.apache.http.params.HttpParams;
import org.apache.http.protocol.BasicHttpContext;
import org.apache.http.protocol.BasicHttpProcessor;
import org.apache.http.protocol.ExecutionContext;
import org.apache.http.protocol.HTTP;
import org.apache.http.protocol.HttpRequestExecutor;
import org.apache.http.protocol.RequestConnControl;
import org.apache.http.protocol.RequestContent;
import org.apache.http.protocol.RequestExpectContinue;
import org.apache.http.protocol.RequestTargetHost;
import org.apache.http.protocol.RequestUserAgent;
import org.apache.http.util.EntityUtils;
import org.apache.log4j.Logger;

import wz.legend.gui.Tray;
import wz.legend.handlers.Handler;
import wz.legend.handlers.RequestHandler;
import wz.legend.handlers.ResponseHandler;
import wz.legend.util.RepeatingHttpEntity;
import wz.legend.util.ScriptStarter;
import wz.legend.util.ScriptStarterImpl;

import java.io.*;
import java.lang.reflect.Array;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.channels.ServerSocketChannel;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

public class ProxyServer extends Thread {
    /**
     * Logger for this class
     */
    private static final Logger LOG = Logger.getLogger(ProxyServer.class);

    public static final int DEFAULT_PORT = 8080;

    private ServerSocket server = null;
    private int thisPort = DEFAULT_PORT;
    private String fwdServer = "";
    private int fwdPort = 0;
    private int ptTimeout = ProxyThread.DEFAULT_TIMEOUT;
    private int debugLevel = 0;

    private List<Handler> handlers = new LinkedList<Handler>();
    private static Context context;

    static {
        context = new Context();
        context.setScriptStater(new ScriptStarterImpl());
    }

    public static Context getContext() {
        return context;
    }

    // private PrintStream debugOut = System.out;

    /* here's a main method, in case you want to run this by itself */
    public static void main(String args[]) {
        int port = 0;
        String fwdProxyServer = "";
        int fwdProxyPort = 0;

        if (args.length == 0) {
            System.err
                    .println("USAGE: java jProxy <port number> [<fwd proxy> <fwd port>]");
            System.err
                    .println("  <port number>   the port this service listens on");
            System.err
                    .println("  <fwd proxy>     optional proxy server to forward requests to");
            System.err
                    .println("  <fwd port>      the port that the optional proxy server is on");
            System.err
                    .println("\nHINT: if you don't want to see all the debug information flying by,");
            System.err
                    .println("you can pipe the output to a file or to 'nul' using \">\". For example:");
            System.err
                    .println("  to send output to the file prox.txt: java jProxy 8080 > prox.txt");
            System.err
                    .println("  to make the output go away: java jProxy 8080 > nul");
            return;
        }

        // get the command-line parameters
        port = Integer.parseInt(args[0]);
        if (args.length > 2) {
            fwdProxyServer = args[1];
            fwdProxyPort = Integer.parseInt(args[2]);
        }

        // create and start the jProxy thread, using a 20 second timeout
        // value to keep the threads from piling up too much
        System.err.println("  **  Starting jProxy on port " + port
                + ". Press CTRL-C to end.  **\n");
        ProxyServer jp = new ProxyServer(port, fwdProxyServer, fwdProxyPort, 20);
        // jp.setDebug(1, System.out); // or set the debug level to 2 for tons
        // of
        // // output
        jp.start();

        // run forever; if you were calling this class from another
        // program and you wanted to stop the jProxy thread at some
        // point, you could write a loop that waits for a certain
        // condition and then calls jProxy.closeSocket() to kill
        // the running jProxy thread
        while (true) {
            try {
                Thread.sleep(3000);
            } catch (Exception e) {
            }
        }

        // if we ever had a condition that stopped the loop above,
        // we'd want to do this to kill the running thread
        // jp.closeSocket();
        // return;
    }

    /**
     * the proxy server just listens for connections and creates a new thread
     * for each connection attempt (the ProxyThread class really does all the
     * work)
     */
    public ProxyServer(int port) {
        this(port, "", 0);
    }

    /**
     * @param port
     * @param proxyServer
     * @param proxyPort
     */
    public ProxyServer(int port, String proxyServer, int proxyPort) {
        this(port, proxyServer, proxyPort, ProxyThread.DEFAULT_TIMEOUT);
    }

    /**
     * @param port
     * @param proxyServer
     * @param proxyPort
     * @param timeout
     */
    public ProxyServer(int port, String proxyServer, int proxyPort, int timeout) {
        thisPort = port;
        fwdServer = proxyServer;
        fwdPort = proxyPort;
        ptTimeout = timeout;

        initHandlers();

        context.setStatus(Status.INACTIVE);

        Tray.getInstance();
    }

    public void initHandlers() {

        try {
            Digester digester = new Digester();
            digester.setValidating(false);
            digester.addObjectCreate("handlers", LinkedList.class);
            digester.addObjectCreate("handlers/handler", "class", Object.class);
            digester.addSetNext("handlers/handler", "add");
            digester.addSetProperties("handlers/handler/param");
            InputStream stream = getClass().getClassLoader()
                    .getResourceAsStream("handlers.xml");
            handlers = (List<Handler>) digester.parse(stream);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // /**
    // * allow the user to decide whether or not to send debug output to the
    // * console or some other PrintStream
    // */
    // public void setDebug(int level, PrintStream out) {
    // debugLevel = level;
    // debugOut = out;
    // }

    public List<Handler> getHandlers() {
        return Collections.unmodifiableList(handlers);
    }

    /**
     * get the port that we're supposed to be listening on
     */
    public int getPort() {
        return thisPort;
    }

    /**
     * return whether or not the socket is currently open
     */
    public boolean isRunning() {
        if (server == null)
            return false;
        else
            return true;
    }

    /**
     * closeSocket will close the open ServerSocket; use this to halt a running
     * jProxy thread
     */
    public void closeSocket() {
        try {
            // close the open server socket
            server.close();
            // send it a message to make it stop waiting immediately
            // (not really necessary)
            /*
             * Socket s = new Socket("localhost", thisPort); OutputStream os =
             * s.getOutputStream(); os.write((byte)0); os.close(); s.close();
             */
        } catch (Exception e) {
            LOG.error(e, e);
        }

        server = null;
    }

    public void run() {
        try {
            // create a server socket, and loop forever listening for
            // client connections
            server = new ServerSocket(thisPort);
            LOG.debug("Started jProxy on port " + thisPort);

            while (true) {
                Socket client = server.accept();
                ProxyThread t = new ProxyThread(client, fwdServer, fwdPort);
                t.setHandlers(getHandlers());
                t.setTimeout(ptTimeout);
                t.start();
            }
        } catch (Exception e) {
            LOG.error("jProxy Thread error: " + e, e);
        }

        closeSocket();
    }

}

/**
 * The ProxyThread will take an HTTP request from the client socket and send it
 * to either the server that the client is trying to contact, or another proxy
 * server
 */
class ProxyThread extends Thread {

    private List<Handler> handlers = new LinkedList<Handler>();

    /**
     * Logger for this class
     */
    private static final Logger LOG = Logger.getLogger(ProxyThread.class);

    private Socket pSocket;
    private String fwdServer = "";
    private int fwdPort = 0;
    // private int debugLevel = 1;

    // the socketTimeout is used to time out the connection to
    // the remote server after a certain period of inactivity;
    // the value is in milliseconds -- use zero if you don't want
    // a timeout
    public static final int DEFAULT_TIMEOUT = 20 * 1000;
    private int socketTimeout = DEFAULT_TIMEOUT;

    public ProxyThread(Socket s) {
        pSocket = s;
    }

    public ProxyThread(Socket s, String proxy, int port) {
        pSocket = s;
        fwdServer = proxy;
        fwdPort = port;
    }

    public void addHandler(Handler handler) {
        handlers.add(handler);
    }

    public void setHandlers(List<Handler> handlers) {
        this.handlers = handlers;
    }

    public void setTimeout(int timeout) {
        // assume that the user will pass the timeout value
        // in seconds (because that's just more intuitive)
        socketTimeout = timeout * 1000;
    }

    public void run() {
        try {
            long startTime = System.currentTimeMillis();

            Socket server = null;

            BasicHttpParams params = new BasicHttpParams();
            DefaultHttpServerConnection serverConnection = new DefaultHttpServerConnection() {
                @Override
                public void bind(final Socket socket, final HttpParams params)
                        throws IOException {
                    super.bind(socket, params);
                }
            };

            serverConnection.bind(pSocket, params);
            HttpRequest httpRequest = serverConnection.receiveRequestHeader();
            if (httpRequest instanceof HttpEntityEnclosingRequest) {
                serverConnection
                        .receiveRequestEntity((HttpEntityEnclosingRequest) httpRequest);
            }

            if (!processRequest(httpRequest))
                return;

            URI uri = new URI(httpRequest.getRequestLine().getUri());
            int hostPort = uri.getPort() > 0 ? uri.getPort() : 80;

            HttpResponse httpResponse = null;
            // either forward this request to another proxy server or
            // send it straight to the Host
            try {
                if ((fwdServer.length() > 0) && (fwdPort > 0)) {
                    server = new Socket(fwdServer, fwdPort);
                } else {
                    server = new Socket(uri.getHost(), hostPort);
                }

                if (server != null) {
                    BasicHttpProcessor httpproc = new BasicHttpProcessor();
                    // Required protocol interceptors
                    httpproc.addInterceptor(new RequestContent());
                    httpproc.addInterceptor(new RequestTargetHost());
                    // Recommended protocol interceptors
                    httpproc.addInterceptor(new RequestConnControl());
                    httpproc.addInterceptor(new RequestUserAgent());
                    httpproc.addInterceptor(new RequestExpectContinue());

                    HttpRequestExecutor httpexecutor = new HttpRequestExecutor();

                    org.apache.http.protocol.HttpContext context = new BasicHttpContext(
                            null);

                    HttpHost httpHost = new HttpHost(uri.getHost(), hostPort);

                    DefaultHttpClientConnection conn = new DefaultHttpClientConnection();
                    ConnectionReuseStrategy connStrategy = new DefaultConnectionReuseStrategy();

                    context
                            .setAttribute(ExecutionContext.HTTP_CONNECTION,
                                    conn);
                    context.setAttribute(ExecutionContext.HTTP_TARGET_HOST,
                            httpHost);
                    context.setAttribute(ExecutionContext.HTTP_REQUEST,
                            httpRequest);

                    conn.bind(server, params);
                    httpRequest.removeHeaders(HTTP.CONTENT_LEN);
                    httpexecutor.preProcess(httpRequest, httpproc, context);
                    httpResponse = httpexecutor.execute(httpRequest, conn,
                            context);
                    httpResponse.setParams(params);
                    httpexecutor.postProcess(httpResponse, httpproc, context);

                    if (httpResponse.getEntity() != null)
                        httpResponse.setEntity(new RepeatingHttpEntity(
                                httpResponse.getEntity()));
                }
            } catch (Exception e) {
                LOG.error("Error1 in ProxyThread: " + e, e);
                httpResponse = new DefaultHttpResponseFactory().newHttpResponse(HttpVersion.HTTP_1_0, 
                        HttpStatus.SC_INTERNAL_SERVER_ERROR, new BasicHttpContext(null));
            } finally {
                if (server != null)
                    server.close();
            }

            serverConnection.sendResponseHeader(httpResponse);
            serverConnection.sendResponseEntity(httpResponse);
            serverConnection.flush();
            serverConnection.close();

            if (httpResponse != null)
                processResponse(httpRequest, httpResponse);

        } catch (URISyntaxException e) {
        } catch (Exception e) {
            LOG.error("Error in ProxyThread: " + e, e);
        } finally {
            try {
                pSocket.close();
            } catch (IOException e) {
            }
        }

    }

    private boolean processRequest(HttpRequest httpRequest) {
        for (Handler handler : handlers) {
            if (handler instanceof RequestHandler) {
                if (!((RequestHandler) handler).handle(httpRequest, ProxyServer
                        .getContext()))
                    return false;
            }
        }
        return true;
    }

    private void processResponse(HttpRequest request, HttpResponse response) {
        LOG.trace(request.getRequestLine().getUri());

        HttpContext httpContext = new HttpContext(request, response);
        for (Handler handler : handlers) {
            try {
                if (handler instanceof ResponseHandler) {
                    ((ResponseHandler) handler).handle(httpContext, ProxyServer
                            .getContext());
                }
            } catch (Throwable e) {
                LOG.error(e, e);
            }
        }
        LOG.trace(request.getRequestLine().getUri() + " end");
    }
}
